# -*- coding: UTF-8 -*-
"""
@项目名称：is_protein.py
@作   者：陆地起飞全靠浪
@创建日期：2025-10-19-15:31
"""
from rdkit import Chem


def is_protein(pdb_path):
    """
    判断PDB文件是蛋白质还是化合物（小分子）
    返回True：蛋白质；返回False：化合物
    """
    try:
        # 读取PDB文件（不忽略H原子，避免影响原子计数）
        mol = Chem.MolFromPDBFile(pdb_path, removeHs=False)
        if mol is None:
            raise ValueError("无法解析PDB文件")

        # 1. 原子数量判断：蛋白质原子数通常远大于小分子（阈值可根据需求调整）
        # num_atoms = mol.GetNumAtoms()
        # if num_atoms < 500:  # 小分子一般小于500原子（粗略阈值）
        #     return False

        # 2. 残基类型判断：检查是否包含典型氨基酸残基
        # 获取所有残基的名称（RDKit中通过PDBResidueInfo获取）
        residues = []
        for atom in mol.GetAtoms():
            pdb_info = atom.GetPDBResidueInfo()
            if pdb_info:
                res_name = pdb_info.GetResidueName().strip()
                if 'UNK' != res_name:
                    residues.append(res_name)

        # # 典型氨基酸残基列表（20种常见氨基酸的三字母代码）
        # amino_acids = {"ALA", "ARG", "ASN", "ASP", "CYS", "GLN", "GLU", "GLY", "HIS", "ILE",
        #                "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL"}
        #
        # # 若存在多个氨基酸残基，则大概率是蛋白质
        # if len(set(residues) & amino_acids) >= 5:  # 至少5种不同氨基酸残基
        #     return True
        if len(residues) > 1:
            print('residues：', residues)
            return True

        # 3. 链结构判断：蛋白质通常有多条链
        chains = set()
        for atom in mol.GetAtoms():
            pdb_info = atom.GetPDBResidueInfo()
            if pdb_info:
                chain_id = pdb_info.GetChainId().strip()
                if chain_id:
                    chains.add(chain_id)
        if len(chains) >= 2:  # 多条链更可能是蛋白质
            return True

        # 若以上条件均不满足，判断为小分子
        return False

    except Exception as e:
        print(f"错误：{e}")
        return None


if __name__ == '__main__':
    # 示例用法
    pdb_file = "/4T/AIDDserver/DataSrc/测试用例/Y001/extracted_chains_class/chain_配体_1.pdb"
    result = is_protein(pdb_file)
    if result is True:
        print(f"{pdb_file} 是蛋白质")
    elif result is False:
        print(f"{pdb_file} 是化合物（小分子）")
    else:
        print("无法判断")
